import { defineStore } from 'pinia'
import axios from 'axios'

interface ImageData {
    id: string
    url: string
    title: string
    prompt: string
    user: {
        username: string
        avatar: string
    }
}

interface GenerateOptions {
    negative_prompt?: string
    model?: string
    sampler?: string
    steps?: number
    width?: number
    height?: number
    init_image?: string
    strength?: number
    mask?: string
}

// ComfyUI 工作流构建器
class WorkflowBuilder {
    private nodeId = 1
    private nodes: Record<string, any> = {}

    private addNode(className: string, inputs: Record<string, any> = {}) {
        const id = this.nodeId++
        this.nodes[id] = {
            class_type: className,
            inputs
        }
        return id
    }

    buildTextToImage(prompt: string, options: GenerateOptions) {
        // 加载模型
        const checkpointId = this.addNode('CheckpointLoaderSimple', {
            ckpt_name: options.model || 'v1-5-pruned.ckpt'
        })

        // 编码正向提示词
        const positiveId = this.addNode('CLIPTextEncode', {
            text: prompt,
            clip: [checkpointId, 0]
        })

        // 编码负向提示词
        const negativeId = this.addNode('CLIPTextEncode', {
            text: options.negative_prompt || '',
            clip: [checkpointId, 0]
        })

        // 设置采样器
        const sampleId = this.addNode('KSampler', {
            model: [checkpointId, 0],
            positive: [positiveId, 0],
            negative: [negativeId, 0],
            sampler_name: options.sampler || 'euler_a',
            steps: options.steps || 20,
            cfg: 7,
            seed: Math.floor(Math.random() * 1000000),
            width: options.width || 512,
            height: options.height || 512
        })

        // 保存图片
        this.addNode('SaveImage', {
            images: [sampleId, 0],
            filename_prefix: 'generated'
        })

        return this.nodes
    }

    buildImageToImage(prompt: string, options: GenerateOptions) {
        // 加载模型
        const checkpointId = this.addNode('CheckpointLoaderSimple', {
            ckpt_name: options.model || 'v1-5-pruned.ckpt'
        })

        // 加载原图
        const loadImageId = this.addNode('LoadImage', {
            image: options.init_image
        })

        // VAE 编码
        const vaeEncodeId = this.addNode('VAEEncode', {
            pixels: [loadImageId, 0],
            vae: [checkpointId, 2]
        })

        // 编码提示词
        const positiveId = this.addNode('CLIPTextEncode', {
            text: prompt,
            clip: [checkpointId, 0]
        })

        const negativeId = this.addNode('CLIPTextEncode', {
            text: options.negative_prompt || '',
            clip: [checkpointId, 0]
        })

        // 采样
        const sampleId = this.addNode('KSampler', {
            model: [checkpointId, 0],
            positive: [positiveId, 0],
            negative: [negativeId, 0],
            latent_image: [vaeEncodeId, 0],
            sampler_name: options.sampler || 'euler_a',
            steps: options.steps || 20,
            cfg: 7,
            seed: Math.floor(Math.random() * 1000000),
            denoise: options.strength || 0.75
        })

        // VAE 解码
        const vaeDecodeId = this.addNode('VAEDecode', {
            samples: [sampleId, 0],
            vae: [checkpointId, 2]
        })

        // 保存图片
        this.addNode('SaveImage', {
            images: [vaeDecodeId, 0],
            filename_prefix: 'generated'
        })

        return this.nodes
    }

    buildInpaint(prompt: string, options: GenerateOptions) {
        // 加载模型
        const checkpointId = this.addNode('CheckpointLoaderSimple', {
            ckpt_name: options.model || 'v1-5-pruned.ckpt'
        })

        // 加载原图和遮罩
        const loadImageId = this.addNode('LoadImage', {
            image: options.init_image
        })

        const loadMaskId = this.addNode('LoadImage', {
            image: options.mask
        })

        // VAE 编码
        const vaeEncodeId = this.addNode('VAEEncode', {
            pixels: [loadImageId, 0],
            vae: [checkpointId, 2]
        })

        // 编码提示词
        const positiveId = this.addNode('CLIPTextEncode', {
            text: prompt,
            clip: [checkpointId, 0]
        })

        const negativeId = this.addNode('CLIPTextEncode', {
            text: options.negative_prompt || '',
            clip: [checkpointId, 0]
        })

        // 设置遮罩区域
        const maskId = this.addNode('SetLatentNoiseMask', {
            samples: [vaeEncodeId, 0],
            mask: [loadMaskId, 0]
        })

        // 采样
        const sampleId = this.addNode('KSampler', {
            model: [checkpointId, 0],
            positive: [positiveId, 0],
            negative: [negativeId, 0],
            latent_image: [maskId, 0],
            sampler_name: options.sampler || 'euler_a',
            steps: options.steps || 20,
            cfg: 7,
            seed: Math.floor(Math.random() * 1000000),
            denoise: options.strength || 1.0
        })

        // VAE 解码
        const vaeDecodeId = this.addNode('VAEDecode', {
            samples: [sampleId, 0],
            vae: [checkpointId, 2]
        })

        // 保存图片
        this.addNode('SaveImage', {
            images: [vaeDecodeId, 0],
            filename_prefix: 'generated'
        })

        return this.nodes
    }
}

export const useImageStore = defineStore('image', {
    state: () => ({
        images: [] as ImageData[],
        loading: false,
        error: null as string | null
    }),

    actions: {
        async generateImage(prompt: string, options: GenerateOptions = {}) {
            this.loading = true
            this.error = null

            try {
                const builder = new WorkflowBuilder()
                let workflow

                if (options.mask) {
                    // 局部重绘模式
                    workflow = builder.buildInpaint(prompt, options)
                } else if (options.init_image) {
                    // 图生图模式
                    workflow = builder.buildImageToImage(prompt, options)
                } else {
                    // 文生图模式
                    workflow = builder.buildTextToImage(prompt, options)
                }

                // 发送请求到 ComfyUI
                const response = await axios.post(
                    `${import.meta.env.VITE_COMFYUI_URL}/prompt`,
                    { prompt: workflow }
                )

                // 等待生成完成
                const result = await this.waitForResult(response.data.prompt_id)

                // 保存生成记录
                const imageData = {
                    title: prompt.slice(0, 50),
                    prompt,
                    url: result.imageUrl
                }
                await this.saveImage(imageData)

                return result
            } catch (err: any) {
                this.error = err.message
                console.error('Generation failed:', err)
                throw err
            } finally {
                this.loading = false
            }
        },

        async waitForResult(promptId: string) {
            const maxRetries = 60
            const interval = 1000

            for (let i = 0; i < maxRetries; i++) {
                try {
                    const response = await axios.get(
                        `${import.meta.env.VITE_COMFYUI_URL}/history/${promptId}`
                    )

                    if (response.data.completed) {
                        const outputs = response.data.outputs || {}
                        const node = Object.values(outputs)[0] as any
                        if (node?.images?.[0]) {
                            const filename = node.images[0].filename
                            return {
                                imageUrl: `${import.meta.env.VITE_COMFYUI_URL}/view?filename=${filename}`,
                                filename
                            }
                        }
                    }
                } catch (err) {
                    console.warn('Failed to check result:', err)
                }

                await new Promise(resolve => setTimeout(resolve, interval))
            }

            throw new Error('Generation timeout')
        }
    }
}) 